package dev.alm.ragdemo.config;

import dev.alm.ragdemo.service.ChatService;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.store.embedding.qdrant.QdrantEmbeddingStore;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.QdrantGrpcClient;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class LLMConfig {

    @Value("${llm.api.key}")
    private String apiKey;

    @Value("${llm.base.url}")
    private String baseUrl;

    @Value("${llm.model.name}")
    private String modelName;

    @Value("${embedding.model.api.key}")
    private String embeddingApiKey;

    @Value("${embedding.model.base.url}")
    private String embeddingBaseUrl;

    @Value("${embedding.model.name}")
    private String embeddingModelName;

    @Bean
    public OpenAiEmbeddingModel openAiEmbeddingModel() {
        return OpenAiEmbeddingModel.builder()
                .apiKey(embeddingApiKey)
                .baseUrl(embeddingBaseUrl)
                .modelName(embeddingModelName)
                .build();
    }

    // 本次使用InMemoryEmbeddingStore
    @Bean
    public InMemoryEmbeddingStore<TextSegment> embeddingStore() {
        return new InMemoryEmbeddingStore<>();
    }

//    // Qdrant客户端
//    @Bean
//    public QdrantClient qdrantClient() {
//        QdrantGrpcClient.Builder builder = QdrantGrpcClient.newBuilder("127.0.0.1", 6334, false);
//        QdrantGrpcClient grpcClient = builder.build();
//        QdrantClient qdrantClient = new QdrantClient(grpcClient);
//        return qdrantClient;
//    }
//
//    @Bean
//    public EmbeddingStore<TextSegment> embeddingStore() {
//        QdrantEmbeddingStore store = QdrantEmbeddingStore.builder()
//                .host("127.0.0.1")
//                .port(6334)
//                .collectionName("test_qdrant")
//                .build();
//        return store;
//    }

    @Bean
    public OpenAiChatModel openAiChatModel() {
        return OpenAiChatModel.builder()
                .apiKey(apiKey)
                .baseUrl(baseUrl)
                .modelName(modelName)
                .build();
    }

    @Bean
    public OpenAiStreamingChatModel openAiStreamingChatModel() {
        return OpenAiStreamingChatModel.builder()
                .apiKey(apiKey)
                .baseUrl(baseUrl)
                .modelName(modelName)
                .build();
    }

    @Bean
    public ChatService chatService(OpenAiChatModel openAiChatModel, OpenAiStreamingChatModel openAiStreamingChatModel, EmbeddingStore<TextSegment> embeddingStore) {
        ChatService service = AiServices.builder(ChatService.class)
                .chatModel(openAiChatModel)
                .streamingChatModel(openAiStreamingChatModel)
                .contentRetriever(EmbeddingStoreContentRetriever.from(embeddingStore))
                .build();
        return service;
    }

}
